"""
ACME 客户端实现
支持 Let's Encrypt 证书申请
"""
import asyncio
import base64
import hashlib
import json
import logging
from typing import List, Optional, Dict, Tuple
from datetime import datetime, timedelta

# ACME 相关库
from acme import client, messages, challenges, errors
import josepy as jose
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, ec
from cryptography.hazmat.backends import default_backend
from cryptography.x509.oid import NameOID
import httpx

# 日志配置
logger = logging.getLogger('acme_client')
logger.setLevel(logging.INFO)
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter('%(levelname)s: %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)


# ========================================================================================
# 配置常量
# ========================================================================================
class ACMEConfig:
    """ACME 配置"""
    
    # Let's Encrypt 目录 URL
    LETSENCRYPT_PROD = "https://acme-v02.api.letsencrypt.org/directory"
    LETSENCRYPT_STAGING = "https://acme-staging-v02.api.letsencrypt.org/directory"
    
    # ZeroSSL 目录 URL
    ZEROSSL_PROD = "https://acme.zerossl.com/v2/DV90"
    
    # 验证超时设置
    CHALLENGE_VERIFICATION_TIMEOUT = 60  # 秒
    CHALLENGE_STATUS_CHECK_INTERVAL = 3  # 秒
    CHALLENGE_MAX_ATTEMPTS = 20
    
    ORDER_STATUS_CHECK_INTERVAL = 2  # 秒
    ORDER_READY_MAX_ATTEMPTS = 30


# ========================================================================================
# 工具函数
# ========================================================================================
def base64url_encode(data: bytes) -> str:
    """Base64 URL 安全编码"""
    encoded = base64.urlsafe_b64encode(data)
    return encoded.rstrip(b'=').decode('ascii')


def base64url_decode(input_str: str) -> bytes:
    """Base64 URL 安全解码"""
    input_str = input_str.replace('-', '+').replace('_', '/')
    padding = 4 - (len(input_str) % 4)
    if padding != 4:
        input_str += '=' * padding
    return base64.b64decode(input_str)


def compute_dns_value(token: str, thumbprint: str) -> str:
    """计算 DNS TXT 记录值"""
    key_authorization = f"{token}.{thumbprint}"
    digest = hashlib.sha256(key_authorization.encode('utf-8')).digest()
    return base64url_encode(digest)


def generate_account_key(algorithm: str = "RSA") -> jose.JWK:
    """生成账户密钥"""
    if algorithm.upper() == "RSA":
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        return jose.JWKRSA(key=private_key)
    elif algorithm.upper() == "ECC":
        private_key = ec.generate_private_key(
            ec.SECP256R1(), backend=default_backend()
        )
        return jose.JWKEC(key=private_key)
    else:
        raise ValueError(f"不支持的算法: {algorithm}")


def generate_certificate_key(algorithm: str = "RSA"):
    """生成证书私钥"""
    if algorithm.upper() == "RSA":
        return rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
    elif algorithm.upper() == "ECC":
        return ec.generate_private_key(
            ec.SECP256R1(), backend=default_backend()
        )
    else:
        raise ValueError(f"不支持的算法: {algorithm}")


def create_csr(domains: List[str], private_key) -> bytes:
    """创建证书签名请求 (CSR)"""
    # 主域名
    common_name = domains[0]
    
    # 构建 Subject
    subject = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, common_name)
    ])
    
    # 构建 CSR
    builder = x509.CertificateSigningRequestBuilder()
    builder = builder.subject_name(subject)
    
    # 添加 SAN（Subject Alternative Names）
    san_list = [x509.DNSName(domain) for domain in domains]
    builder = builder.add_extension(
        x509.SubjectAlternativeName(san_list),
        critical=False
    )
    
    # 签名 CSR
    csr = builder.sign(private_key, hashes.SHA256(), backend=default_backend())
    
    return csr.public_bytes(serialization.Encoding.DER)


# ========================================================================================
# ACME 客户端
# ========================================================================================
class ACMEClient:
    """ACME 客户端"""
    
    def __init__(
        self,
        ca_provider: str = "letsencrypt",
        contact_email: str = None,
        account_key: Optional[jose.JWK] = None
    ):
        """
        初始化 ACME 客户端
        
        Args:
            ca_provider: CA 提供商 (letsencrypt, zerossl)
            contact_email: 联系邮箱
            account_key: 账户密钥（可选，不提供则自动生成）
        """
        self.ca_provider = ca_provider.lower()
        self.contact_email = contact_email
        
        # 选择目录 URL
        if self.ca_provider == "letsencrypt":
            self.directory_url = ACMEConfig.LETSENCRYPT_PROD
        elif self.ca_provider == "zerossl":
            self.directory_url = ACMEConfig.ZEROSSL_PROD
        else:
            raise ValueError(f"不支持的 CA 提供商: {ca_provider}")
        
        # 账户密钥
        if account_key:
            self.account_key = account_key
        else:
            logger.info("生成新的账户密钥...")
            self.account_key = generate_account_key("RSA")
        
        # ACME 客户端
        self.net = client.ClientNetwork(self.account_key, user_agent="SAC/1.0")
        self.directory = None
        self.client = None
        self.account = None
    
    async def initialize(self):
        """初始化客户端（连接到 ACME 服务器）"""
        logger.info(f"连接到 ACME 服务器: {self.directory_url}")
        
        # 获取目录
        self.directory = messages.Directory.from_json(
            self.net.get(self.directory_url).json()
        )
        self.client = client.ClientV2(self.directory, self.net)
        
        # 注册或查找账户
        if self.contact_email:
            email = f"mailto:{self.contact_email}"
        else:
            email = None
        
        try:
            # 尝试注册新账户
            self.account = self.client.new_account(
                messages.NewRegistration.from_data(
                    email=email,
                    terms_of_service_agreed=True
                )
            )
            logger.info("✅ 账户注册成功")
        except errors.ConflictError:
            # 账户已存在
            logger.info("账户已存在，使用现有账户")
            self.account = self.client.new_account(
                messages.NewRegistration.from_data(
                    email=email,
                    terms_of_service_agreed=True,
                    only_return_existing=True
                )
            )
        
        logger.info(f"账户 URI: {self.account.uri}")
        return self
    
    def get_jwk_thumbprint(self) -> str:
        """获取 JWK 指纹"""
        return base64url_encode(self.account_key.thumbprint())
    
    async def create_order(self, domains: List[str]) -> Tuple[messages.OrderResource, List[Dict]]:
        """
        创建证书订单
        
        Args:
            domains: 域名列表
        
        Returns:
            (order, dns_records): 订单资源和 DNS 验证记录
        """
        logger.info(f"创建证书订单，域名: {', '.join(domains)}")
        
        # 创建订单
        order = self.client.new_order(
            csr_pem=None  # 稍后提供
        )
        
        # 解析授权
        dns_records = []
        for authz in order.authorizations:
            domain = authz.body.identifier.value
            logger.info(f"处理域名授权: {domain}")
            
            # 查找 DNS-01 挑战
            dns_challenge = None
            for challenge in authz.body.challenges:
                if isinstance(challenge.chall, challenges.DNS01):
                    dns_challenge = challenge
                    break
            
            if not dns_challenge:
                raise ValueError(f"域名 {domain} 没有 DNS-01 挑战")
            
            # 计算 DNS TXT 记录值
            token = dns_challenge.chall.token
            thumbprint = self.get_jwk_thumbprint()
            dns_value = compute_dns_value(token, thumbprint)
            
            dns_records.append({
                "domain": domain,
                "record_name": f"_acme-challenge.{domain}",
                "record_type": "TXT",
                "record_value": dns_value,
                "token": token,
                "challenge_uri": dns_challenge.uri,
                "validated": False
            })
            
            logger.info(f"DNS 记录: _acme-challenge.{domain} -> {dns_value}")
        
        return order, dns_records
    
    async def verify_dns_challenge(
        self,
        order: messages.OrderResource,
        dns_records: List[Dict]
    ) -> bool:
        """
        验证 DNS 挑战
        
        Args:
            order: 订单资源
            dns_records: DNS 记录列表
        
        Returns:
            bool: 是否全部验证成功
        """
        logger.info("开始验证 DNS 挑战...")
        
        # 等待 DNS 记录生效
        logger.info("等待 DNS 记录生效...")
        await asyncio.sleep(10)
        
        # 通知 ACME 服务器进行验证
        for authz in order.authorizations:
            domain = authz.body.identifier.value
            
            # 查找对应的 DNS-01 挑战
            dns_challenge = None
            for challenge in authz.body.challenges:
                if isinstance(challenge.chall, challenges.DNS01):
                    dns_challenge = challenge
                    break
            
            if not dns_challenge:
                continue
            
            logger.info(f"通知验证: {domain}")
            
            # 触发验证
            response = self.client.answer_challenge(
                dns_challenge,
                dns_challenge.response(self.account_key)
            )
            
            # 轮询验证状态
            for attempt in range(ACMEConfig.CHALLENGE_MAX_ATTEMPTS):
                await asyncio.sleep(ACMEConfig.CHALLENGE_STATUS_CHECK_INTERVAL)
                
                challenge_status = self.client.poll(dns_challenge)
                status = challenge_status.body.status
                
                logger.info(f"验证状态 [{domain}]: {status}")
                
                if status == messages.STATUS_VALID:
                    logger.info(f"✅ 域名 {domain} 验证成功")
                    break
                elif status == messages.STATUS_INVALID:
                    error_detail = challenge_status.body.error
                    raise ValueError(f"域名 {domain} 验证失败: {error_detail}")
                elif status == messages.STATUS_PENDING:
                    continue
                else:
                    logger.warning(f"未知状态: {status}")
            else:
                raise TimeoutError(f"域名 {domain} 验证超时")
        
        logger.info("✅ 所有域名验证成功")
        return True
    
    async def finalize_order(
        self,
        order: messages.OrderResource,
        domains: List[str],
        algorithm: str = "RSA"
    ) -> Tuple[str, str, datetime]:
        """
        完成订单并获取证书
        
        Args:
            order: 订单资源
            domains: 域名列表
            algorithm: 证书算法 (RSA/ECC)
        
        Returns:
            (cert_pem, key_pem, expires_at): 证书、私钥、过期时间
        """
        logger.info("生成证书私钥和 CSR...")
        
        # 生成证书私钥
        cert_key = generate_certificate_key(algorithm)
        
        # 创建 CSR
        csr_der = create_csr(domains, cert_key)
        
        # 完成订单
        logger.info("提交 CSR...")
        order = self.client.finalize_order(order, csr_der)
        
        # 等待证书就绪
        logger.info("等待证书签发...")
        for attempt in range(ACMEConfig.ORDER_READY_MAX_ATTEMPTS):
            await asyncio.sleep(ACMEConfig.ORDER_STATUS_CHECK_INTERVAL)
            
            order = self.client.poll_and_finalize(order)
            
            if order.fullchain_pem:
                logger.info("✅ 证书签发成功")
                break
        else:
            raise TimeoutError("证书签发超时")
        
        # 导出证书
        cert_pem = order.fullchain_pem
        
        # 导出私钥
        key_pem = cert_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption()
        ).decode('utf-8')
        
        # 解析证书获取过期时间
        cert_obj = x509.load_pem_x509_certificate(
            cert_pem.encode('utf-8'),
            backend=default_backend()
        )
        expires_at = cert_obj.not_valid_after
        
        logger.info(f"证书过期时间: {expires_at}")
        
        return cert_pem, key_pem, expires_at
    
    async def issue_certificate(
        self,
        domains: List[str],
        algorithm: str = "RSA",
        dns_records_callback=None
    ) -> Dict:
        """
        申请证书（完整流程）
        
        Args:
            domains: 域名列表
            algorithm: 证书算法
            dns_records_callback: DNS 记录回调函数 async (records) -> None
        
        Returns:
            证书信息字典
        """
        try:
            # 1. 初始化客户端
            await self.initialize()
            
            # 2. 创建订单
            order, dns_records = await self.create_order(domains)
            
            # 3. 如果有回调，调用回调函数处理 DNS 记录
            if dns_records_callback:
                await dns_records_callback(dns_records)
            else:
                # 返回 DNS 记录，等待手动添加
                return {
                    "status": "pending_dns",
                    "dns_records": dns_records,
                    "order": order,
                    "message": "请添加 DNS TXT 记录后调用验证接口"
                }
            
            # 4. 验证 DNS 挑战
            await self.verify_dns_challenge(order, dns_records)
            
            # 5. 完成订单并获取证书
            cert_pem, key_pem, expires_at = await self.finalize_order(
                order, domains, algorithm
            )
            
            return {
                "status": "issued",
                "cert_fullchain_pem": cert_pem,
                "private_key_pem": key_pem,
                "expires_at": expires_at,
                "domains": domains,
                "algorithm": algorithm
            }
            
        except Exception as e:
            logger.error(f"证书申请失败: {str(e)}")
            raise


# ========================================================================================
# DNS 记录验证工具
# ========================================================================================
async def verify_dns_txt_record(domain: str, expected_value: str, timeout: int = 60) -> bool:
    """
    验证 DNS TXT 记录是否已生效
    
    Args:
        domain: 记录名称 (_acme-challenge.example.com)
        expected_value: 期望的 TXT 值
        timeout: 超时时间（秒）
    
    Returns:
        bool: 是否验证成功
    """
    import dns.resolver
    
    logger.info(f"验证 DNS TXT 记录: {domain}")
    
    start_time = datetime.now()
    while (datetime.now() - start_time).seconds < timeout:
        try:
            answers = dns.resolver.resolve(domain, 'TXT')
            for rdata in answers:
                txt_value = rdata.to_text().strip('"')
                if txt_value == expected_value:
                    logger.info(f"✅ DNS TXT 记录已生效: {domain}")
                    return True
            
            logger.info(f"DNS TXT 记录值不匹配，等待...")
            await asyncio.sleep(5)
        except Exception as e:
            logger.warning(f"DNS 查询失败: {str(e)}，继续等待...")
            await asyncio.sleep(5)
    
    logger.error(f"❌ DNS TXT 记录验证超时: {domain}")
    return False

